{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "BERT SUMM.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/santhoshkolloju/bert_summ/blob/master/BERT_SUMM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "pbgEu1oyhPca",
        "colab_type": "code",
        "outputId": "1ecf3d0e-75c4-4f0d-8a64-18fc6bcfb4a1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "device_name = tf.test.gpu_device_name()\n",
        "if device_name != '/device:GPU:0':\n",
        "  raise SystemError('GPU device not found')\n",
        "print('Found GPU at: {}'.format(device_name))"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Found GPU at: /device:GPU:0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "UqeTD_eGkxVd",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from google.colab import auth\n",
        "auth.authenticate_user()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "SFCZikuuibh7",
        "colab_type": "code",
        "outputId": "f228c486-ff0b-4740-c9c2-10a2d9b893ba",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 267
        }
      },
      "cell_type": "code",
      "source": [
        "import sys\n",
        "\n",
        "!test -d texar_repo || git clone https://github.com/asyml/texar.git texar_repo\n",
        "if not 'texar_repo' in sys.path:\n",
        "  sys.path += ['texar_repo']\n",
        "!pip install funcsigs"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'texar_repo'...\n",
            "remote: Enumerating objects: 10204, done.\u001b[K\n",
            "remote: Total 10204 (delta 0), reused 0 (delta 0), pack-reused 10204\u001b[K\n",
            "Receiving objects: 100% (10204/10204), 2.30 MiB | 12.21 MiB/s, done.\n",
            "Resolving deltas: 100% (7774/7774), done.\n",
            "Collecting funcsigs\n",
            "  Downloading https://files.pythonhosted.org/packages/69/cb/f5be453359271714c01b9bd06126eaf2e368f1fddfff30818754b5ac2328/funcsigs-1.0.2-py2.py3-none-any.whl\n",
            "Installing collected packages: funcsigs\n",
            "Successfully installed funcsigs-1.0.2\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "_IB17muzsBd0",
        "colab_type": "code",
        "outputId": "34a14a15-84b6-47d0-d286-2ae052c858d0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 74
        }
      },
      "cell_type": "code",
      "source": [
        "\"\"\"\n",
        "from google_drive_downloader import GoogleDriveDownloader as gdd\n",
        "\n",
        "gdd.download_file_from_google_drive(file_id='0BwmD_VLjROrfTHk4NFg2SndKcjQ',\n",
        "                                    dest_path='./cnn_stories.tgz',\n",
        "                                    unzip=True)\n",
        "!tar -zxf cnn_stories.tgz\n",
        "\n",
        "from os import listdir\n",
        "import string\n",
        "\n",
        "# load doc into memory\n",
        "def load_doc(filename):\n",
        "\t# open the file as read only\n",
        "\tfile = open(filename, encoding='utf-8')\n",
        "\t# read all text\n",
        "\ttext = file.read()\n",
        "\t# close the file\n",
        "\tfile.close()\n",
        "\treturn text\n",
        "\n",
        "# split a document into news story and highlights\n",
        "def split_story(doc):\n",
        "\t# find first highlight\n",
        "\tindex = doc.find('@highlight')\n",
        "\t# split into story and highlights\n",
        "\tstory, highlights = doc[:index], doc[index:].split('@highlight')\n",
        "\t# strip extra white space around each highlight\n",
        "\thighlights = [h.strip() for h in highlights if len(h) > 0]\n",
        "\treturn story, highlights\n",
        "\n",
        "# load all stories in a directory\n",
        "def load_stories(directory):\n",
        "\tstories = list()\n",
        "\tfor name in listdir(directory):\n",
        "\t\tfilename = directory + '/' + name\n",
        "\t\t# load document\n",
        "\t\tdoc = load_doc(filename)\n",
        "\t\t# split into story and highlights\n",
        "\t\tstory, highlights = split_story(doc)\n",
        "\t\t# store\n",
        "\t\tstories.append({'story':story, 'highlights':highlights})\n",
        "\treturn stories\n",
        "\n",
        "# clean a list of lines\n",
        "def clean_lines(lines):\n",
        "\tcleaned = list()\n",
        "\t# prepare a translation table to remove punctuation\n",
        "\ttable = str.maketrans('', '', string.punctuation)\n",
        "\tfor line in lines:\n",
        "\t\t# strip source cnn office if it exists\n",
        "\t\tindex = line.find('(CNN) -- ')\n",
        "\t\tif index > -1:\n",
        "\t\t\tline = line[index+len('(CNN)'):]\n",
        "\t\t# tokenize on white space\n",
        "\t\tline = line.split()\n",
        "\t\t# convert to lower case\n",
        "\t\tline = [word.lower() for word in line]\n",
        "\t\t# remove punctuation from each token\n",
        "\t\tline = [w.translate(table) for w in line]\n",
        "\t\t# remove tokens with numbers in them\n",
        "\t\tline = [word for word in line if word.isalpha()]\n",
        "\t\t# store as string\n",
        "\t\tcleaned.append(' '.join(line))\n",
        "\t# remove empty strings\n",
        "\tcleaned = [c for c in cleaned if len(c) > 0]\n",
        "\treturn cleaned\n",
        "\n",
        "# load stories\n",
        "directory = 'cnn/stories/'\n",
        "stories = load_stories(directory)\n",
        "print('Loaded Stories %d' % len(stories))\n",
        "\n",
        "# clean stories\n",
        "f1 = open(\"stories.txt\",'w')\n",
        "f2 = open(\"summary.txt\",'w')\n",
        "for example in stories:\n",
        "  example['story'] = clean_lines(example['story'].split('\\n'))\n",
        "  example['highlights'] = clean_lines(example['highlights'])\n",
        "  f1.write(\" \".join(example['story']))\n",
        "  f1.write(\"\\n\")\n",
        "  f2.write(\" \".join(example['highlights']))\n",
        "  f2.write(\"\\n\")\n",
        "f1.close()\n",
        "f2.close()\n",
        "  \n",
        "story = open(\"stories.txt\").readlines()\n",
        "summ = open(\"summary.txt\").readlines() \n",
        "train_story = story[0:90000]\n",
        "train_summ = summ[0:90000]\n",
        "\n",
        "eval_story = story[90000:91579]\n",
        "eval_summ = summ[90000:91579]\n",
        "\n",
        "\n",
        "test_story = story[91579:92579]\n",
        "test_summ = summ[91579:92579]\n",
        "\n",
        "\n",
        "with open(\"train_story.txt\",'w') as f:\n",
        "  f.write(\"\\n\".join(train_story))\n",
        "  \n",
        "with open(\"train_summ.txt\",'w') as f:\n",
        "  f.write(\"\\n\".join(train_summ))\n",
        "  \n",
        "with open(\"eval_story.txt\",'w') as f:\n",
        "  f.write(\"\\n\".join(eval_story))\n",
        "  \n",
        "  \n",
        "with open(\"eval_summ.txt\",'w') as f:\n",
        "  f.write(\"\\n\".join(eval_summ))\n",
        "  \n",
        "  \n",
        "with open(\"test_story.txt\",'w') as f:\n",
        "  f.write(\"\\n\".join(test_story))\n",
        "  \n",
        "  \n",
        "with open(\"test_summ.txt\",'w') as f:\n",
        "  f.write(\"\\n\".join(test_summ))  \n",
        "  \"\"\""
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'\\nfrom google_drive_downloader import GoogleDriveDownloader as gdd\\n\\ngdd.download_file_from_google_drive(file_id=\\'0BwmD_VLjROrfTHk4NFg2SndKcjQ\\',\\n                                    dest_path=\\'./cnn_stories.tgz\\',\\n                                    unzip=True)\\n!tar -zxf cnn_stories.tgz\\n\\nfrom os import listdir\\nimport string\\n\\n# load doc into memory\\ndef load_doc(filename):\\n\\t# open the file as read only\\n\\tfile = open(filename, encoding=\\'utf-8\\')\\n\\t# read all text\\n\\ttext = file.read()\\n\\t# close the file\\n\\tfile.close()\\n\\treturn text\\n\\n# split a document into news story and highlights\\ndef split_story(doc):\\n\\t# find first highlight\\n\\tindex = doc.find(\\'@highlight\\')\\n\\t# split into story and highlights\\n\\tstory, highlights = doc[:index], doc[index:].split(\\'@highlight\\')\\n\\t# strip extra white space around each highlight\\n\\thighlights = [h.strip() for h in highlights if len(h) > 0]\\n\\treturn story, highlights\\n\\n# load all stories in a directory\\ndef load_stories(directory):\\n\\tstories = list()\\n\\tfor name in listdir(directory):\\n\\t\\tfilename = directory + \\'/\\' + name\\n\\t\\t# load document\\n\\t\\tdoc = load_doc(filename)\\n\\t\\t# split into story and highlights\\n\\t\\tstory, highlights = split_story(doc)\\n\\t\\t# store\\n\\t\\tstories.append({\\'story\\':story, \\'highlights\\':highlights})\\n\\treturn stories\\n\\n# clean a list of lines\\ndef clean_lines(lines):\\n\\tcleaned = list()\\n\\t# prepare a translation table to remove punctuation\\n\\ttable = str.maketrans(\\'\\', \\'\\', string.punctuation)\\n\\tfor line in lines:\\n\\t\\t# strip source cnn office if it exists\\n\\t\\tindex = line.find(\\'(CNN) -- \\')\\n\\t\\tif index > -1:\\n\\t\\t\\tline = line[index+len(\\'(CNN)\\'):]\\n\\t\\t# tokenize on white space\\n\\t\\tline = line.split()\\n\\t\\t# convert to lower case\\n\\t\\tline = [word.lower() for word in line]\\n\\t\\t# remove punctuation from each token\\n\\t\\tline = [w.translate(table) for w in line]\\n\\t\\t# remove tokens with numbers in them\\n\\t\\tline = [word for word in line if word.isalpha()]\\n\\t\\t# store as string\\n\\t\\tcleaned.append(\\' \\'.join(line))\\n\\t# remove empty strings\\n\\tcleaned = [c for c in cleaned if len(c) > 0]\\n\\treturn cleaned\\n\\n# load stories\\ndirectory = \\'cnn/stories/\\'\\nstories = load_stories(directory)\\nprint(\\'Loaded Stories %d\\' % len(stories))\\n\\n# clean stories\\nf1 = open(\"stories.txt\",\\'w\\')\\nf2 = open(\"summary.txt\",\\'w\\')\\nfor example in stories:\\n  example[\\'story\\'] = clean_lines(example[\\'story\\'].split(\\'\\n\\'))\\n  example[\\'highlights\\'] = clean_lines(example[\\'highlights\\'])\\n  f1.write(\" \".join(example[\\'story\\']))\\n  f1.write(\"\\n\")\\n  f2.write(\" \".join(example[\\'highlights\\']))\\n  f2.write(\"\\n\")\\nf1.close()\\nf2.close()\\n  \\nstory = open(\"stories.txt\").readlines()\\nsumm = open(\"summary.txt\").readlines() \\ntrain_story = story[0:90000]\\ntrain_summ = summ[0:90000]\\n\\neval_story = story[90000:91579]\\neval_summ = summ[90000:91579]\\n\\n\\ntest_story = story[91579:92579]\\ntest_summ = summ[91579:92579]\\n\\n\\nwith open(\"train_story.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(train_story))\\n  \\nwith open(\"train_summ.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(train_summ))\\n  \\nwith open(\"eval_story.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(eval_story))\\n  \\n  \\nwith open(\"eval_summ.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(eval_summ))\\n  \\n  \\nwith open(\"test_story.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(test_story))\\n  \\n  \\nwith open(\"test_summ.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(test_summ))  \\n  '"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 4
        }
      ]
    },
    {
      "metadata": {
        "id": "SsWJmIfmij-_",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import os\n",
        "import csv\n",
        "import collections\n",
        "import sys\n",
        "from texar_repo.examples.bert.utils import data_utils, model_utils, tokenization\n",
        "import importlib\n",
        "import tensorflow as tf\n",
        "import texar as tx \n",
        "from texar_repo.examples.bert import config_classifier as config_downstream\n",
        "from texar_repo.texar.utils import transformer_utils\n",
        "from texar_repo.examples.transformer.utils import data_utils, utils\n",
        "from texar_repo.examples.transformer.bleu_tool import bleu_wrapper\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Ko2McfcdhbcN",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#config\n",
        "\n",
        "dcoder_config = {\n",
        "    'dim': 768,\n",
        "    'num_blocks': 6,\n",
        "    'multihead_attention': {\n",
        "        'num_heads': 8,\n",
        "        'output_dim': 768\n",
        "        # See documentation for more optional hyperparameters\n",
        "    },\n",
        "    'position_embedder_hparams': {\n",
        "        'dim': 768\n",
        "    },\n",
        "    'initializer': {\n",
        "        'type': 'variance_scaling_initializer',\n",
        "        'kwargs': {\n",
        "            'scale': 1.0,\n",
        "            'mode': 'fan_avg',\n",
        "            'distribution': 'uniform',\n",
        "        },\n",
        "    },\n",
        "    'poswise_feedforward': tx.modules.default_transformer_poswise_net_hparams(\n",
        "        output_dim=768)\n",
        "}\n",
        "\n",
        "loss_label_confidence = 0.9\n",
        "\n",
        "random_seed = 1234\n",
        "beam_width = 5\n",
        "alpha = 0.6\n",
        "hidden_dim = 768\n",
        "\n",
        "\n",
        "opt = {\n",
        "    'optimizer': {\n",
        "        'type': 'AdamOptimizer',\n",
        "        'kwargs': {\n",
        "            'beta1': 0.9,\n",
        "            'beta2': 0.997,\n",
        "            'epsilon': 1e-9\n",
        "        }\n",
        "    }\n",
        "}\n",
        "\n",
        "\n",
        "lr = {\n",
        "    'learning_rate_schedule': 'constant.linear_warmup.rsqrt_decay.rsqrt_depth',\n",
        "    'lr_constant': 2 * (hidden_dim ** -0.5),\n",
        "    'static_lr': 1e-3,\n",
        "    'warmup_steps': 2000,\n",
        "}\n",
        "\n",
        "bos_token_id =101\n",
        "eos_token_id = 102\n",
        "\n",
        "model_dir= \"./models\"\n",
        "run_mode= \"train_and_evaluate\"\n",
        "batch_size = 32\n",
        "test_batch_size = 32\n",
        "\n",
        "max_train_epoch = 20\n",
        "display_steps = 100\n",
        "eval_steps = 100000\n",
        "\n",
        "max_decoding_length = 400\n",
        "\n",
        "max_seq_length_src = 512\n",
        "max_seq_length_tgt = 400\n",
        "\n",
        "bert_pretrain_dir = 'bert_pretrained_models/uncased_L-12_H-768_A-12'\n",
        "#config"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "MrBw61rEiXeE",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!mkdir bert_pretrained_models"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "FH2RX773i40g",
        "colab_type": "code",
        "outputId": "3e61c01b-e31b-485b-89e0-b51fddc341d4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 354
        }
      },
      "cell_type": "code",
      "source": [
        "!wget https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip -P bert_pretrained_models/;\n",
        "!unzip bert_pretrained_models/uncased_L-12_H-768_A-12.zip -d bert_pretrained_models/\n"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2019-01-31 01:48:29--  https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip\n",
            "Resolving storage.googleapis.com (storage.googleapis.com)... 64.233.181.128, 2607:f8b0:4001:c0f::80\n",
            "Connecting to storage.googleapis.com (storage.googleapis.com)|64.233.181.128|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 407727028 (389M) [application/zip]\n",
            "Saving to: ‘bert_pretrained_models/uncased_L-12_H-768_A-12.zip’\n",
            "\n",
            "uncased_L-12_H-768_ 100%[===================>] 388.84M   136MB/s    in 2.9s    \n",
            "\n",
            "2019-01-31 01:48:32 (136 MB/s) - ‘bert_pretrained_models/uncased_L-12_H-768_A-12.zip’ saved [407727028/407727028]\n",
            "\n",
            "Archive:  bert_pretrained_models/uncased_L-12_H-768_A-12.zip\n",
            "   creating: bert_pretrained_models/uncased_L-12_H-768_A-12/\n",
            "  inflating: bert_pretrained_models/uncased_L-12_H-768_A-12/bert_model.ckpt.meta  \n",
            "  inflating: bert_pretrained_models/uncased_L-12_H-768_A-12/bert_model.ckpt.data-00000-of-00001  \n",
            "  inflating: bert_pretrained_models/uncased_L-12_H-768_A-12/vocab.txt  \n",
            "  inflating: bert_pretrained_models/uncased_L-12_H-768_A-12/bert_model.ckpt.index  \n",
            "  inflating: bert_pretrained_models/uncased_L-12_H-768_A-12/bert_config.json  \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "JC60vOxGjI-M",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class InputExample():\n",
        "    \"\"\"A single training/test example for simple sequence classification.\"\"\"\n",
        "\n",
        "    def __init__(self, guid, text_a, text_b=None):\n",
        "        \"\"\"Constructs a InputExample.\n",
        "        Args:\n",
        "            guid: Unique id for the example.\n",
        "            text_a: string. The untokenized text of the first sequence.\n",
        "                For single sequence tasks, only this sequence must be specified.\n",
        "            text_b: (Optional) string. The untokenized text of the second\n",
        "                sequence. Only must be specified for sequence pair tasks.\n",
        "            label: (Optional) string. The label of the example. This should be\n",
        "                specified for train and dev examples, but not for test examples.\n",
        "        \"\"\"\n",
        "        self.guid = guid\n",
        "        self.src_txt = text_a\n",
        "        self.tgt_txt = text_b\n",
        "        \n",
        "class InputFeatures():\n",
        "    \"\"\"A single set of features of data.\"\"\"\n",
        "\n",
        "    def __init__(self, src_input_ids,src_input_mask,src_segment_ids,tgt_input_ids,tgt_input_mask,tgt_labels):\n",
        "        self.src_input_ids = src_input_ids\n",
        "        self.src_input_mask = src_input_mask\n",
        "        self.src_segment_ids = src_segment_ids\n",
        "        self.tgt_input_ids = tgt_input_ids\n",
        "        self.tgt_input_mask = tgt_input_mask \n",
        "        self.tgt_labels = tgt_labels\n",
        "        \n",
        "       \n",
        "class DataProcessor(object):\n",
        "    \"\"\"Base class for data converters for sequence classification data sets.\"\"\"\n",
        "\n",
        "    def get_train_examples(self, data_dir):\n",
        "        \"\"\"Gets a collection of `InputExample`s for the train set.\"\"\"\n",
        "        raise NotImplementedError()\n",
        "\n",
        "    def get_dev_examples(self, data_dir):\n",
        "        \"\"\"Gets a collection of `InputExample`s for the dev set.\"\"\"\n",
        "        raise NotImplementedError()\n",
        "\n",
        "    def get_test_examples(self, data_dir):\n",
        "        \"\"\"Gets a collection of `InputExample`s for prediction.\"\"\"\n",
        "        raise NotImplementedError()\n",
        "\n",
        "    def get_labels(self):\n",
        "        \"\"\"Gets the list of labels for this data set.\"\"\"\n",
        "        raise NotImplementedError()\n",
        "\n",
        "    @classmethod\n",
        "    def _read_tsv(cls, input_file, quotechar=None):\n",
        "        \"\"\"Reads a tab separated value file.\"\"\"\n",
        "        with tf.gfile.Open(input_file, \"r\") as f:\n",
        "            reader = csv.reader(f, delimiter=\"\\t\", quotechar=quotechar)\n",
        "            lines = []\n",
        "            i = 0\n",
        "            for line in reader:\n",
        "                lines.append(line)\n",
        "        return lines\n",
        "\n",
        "\n",
        "    @classmethod\n",
        "    def _read_file(cls, input_file, quotechar=None):\n",
        "        \"\"\"Reads a tab separated value file.\"\"\"\n",
        "        with tf.gfile.Open(input_file, \"r\") as f:\n",
        "            reader = csv.reader(f, delimiter=\"\\n\", quotechar=quotechar)\n",
        "            lines = []\n",
        "            i = 0\n",
        "            for line in reader:\n",
        "                lines.append(line)\n",
        "        return lines\n",
        "      \n",
        "      \n",
        "class CNNDailymail(DataProcessor):\n",
        "    \"\"\"Processor for the CoLA data set (GLUE version).\"\"\"\n",
        "\n",
        "    def get_train_examples(self, data_dir):\n",
        "        \"\"\"See base class.\"\"\"\n",
        "        return self._create_examples(\n",
        "            self._read_file(os.path.join(data_dir, \"train_story.txt\")),self._read_file(os.path.join(data_dir, \"train_summ.txt\")),\n",
        "            \"train\")\n",
        "\n",
        "    def get_dev_examples(self, data_dir):\n",
        "        \"\"\"See base class.\"\"\"\n",
        "        return self._create_examples(\n",
        "            self._read_file(os.path.join(data_dir, \"eval_story.txt\")),self._read_file(os.path.join(data_dir, \"eval_summ.txt\")),\n",
        "            \"dev\")\n",
        "\n",
        "    def get_test_examples(self, data_dir):\n",
        "        \"\"\"See base class.\"\"\"\n",
        "        return self._create_examples(\n",
        "            self._read_file(os.path.join(data_dir, \"test_story.txt\")),self._read_file(os.path.join(data_dir, \"test_summ.txt\")),\n",
        "            \"test\")\n",
        "\n",
        "    def _create_examples(self, src_lines,tgt_lines,set_type):\n",
        "        examples = [] \n",
        "        for i,data in enumerate(zip(src_lines,tgt_lines)):\n",
        "            guid = \"%s-%s\" % (set_type, i)\n",
        "            if set_type == \"test\" and i == 0:\n",
        "                continue\n",
        "            else:\n",
        "                #print(data)\n",
        "                if len(data[0])==0 or len(data[1])==0:\n",
        "                  continue\n",
        "                src_lines = tokenization.convert_to_unicode(data[0][0])\n",
        "                tgt_lines = tokenization.convert_to_unicode(data[1][0])\n",
        "                examples.append(InputExample(guid=guid, text_a=src_lines,\n",
        "                                         text_b=tgt_lines))\n",
        "        return examples\n",
        "  \n",
        "  \n",
        "def file_based_convert_examples_to_features(\n",
        "        examples, max_seq_length_src,max_seq_length_tgt,tokenizer, output_file):\n",
        "    \"\"\"Convert a set of `InputExample`s to a TFRecord file.\"\"\"\n",
        "\n",
        "    writer = tf.python_io.TFRecordWriter(output_file)\n",
        "\n",
        "    for (ex_index, example) in enumerate(examples):\n",
        "        #print(\"ex_index\",ex_index)\n",
        "\n",
        "        if (ex_index+1) %1000 == 0 :\n",
        "          print(\"------------processed..{}...examples\".format(ex_index))\n",
        "          \n",
        "        feature = convert_single_example(ex_index, example,\n",
        "                                         max_seq_length_src,max_seq_length_tgt,tokenizer)\n",
        "\n",
        "        def create_int_feature(values):\n",
        "            return tf.train.Feature(\n",
        "                int64_list=tf.train.Int64List(value=list(values)))\n",
        "\n",
        "        features = collections.OrderedDict()\n",
        "        features[\"src_input_ids\"] = create_int_feature(feature.src_input_ids)\n",
        "        features[\"src_input_mask\"] = create_int_feature(feature.src_input_mask)\n",
        "        features[\"src_segment_ids\"] = create_int_feature(feature.src_segment_ids)\n",
        "\n",
        "        features[\"tgt_input_ids\"] = create_int_feature(feature.tgt_input_ids)\n",
        "        features[\"tgt_input_mask\"] = create_int_feature(feature.tgt_input_mask)\n",
        "        features['tgt_labels'] = create_int_feature(feature.tgt_labels)\n",
        "        \n",
        "        \n",
        "        \n",
        "        #print(feature.tgt_labels)\n",
        "        \n",
        "\n",
        "        tf_example = tf.train.Example(\n",
        "            features=tf.train.Features(feature=features))\n",
        "        writer.write(tf_example.SerializeToString())\n",
        "\n",
        "\n",
        "def convert_single_example(ex_index, example, max_seq_length_src,max_seq_length_tgt,\n",
        "                           tokenizer):\n",
        "    \"\"\"Converts a single `InputExample` into a single `InputFeatures`.\"\"\"\n",
        "    \"\"\"\n",
        "    label_map = {}\n",
        "    for (i, label) in enumerate(label_list):\n",
        "        label_map[label] = i\n",
        "    \"\"\"\n",
        "    tokens_a = tokenizer.tokenize(example.src_txt)\n",
        "    tokens_b = tokenizer.tokenize(example.tgt_txt)\n",
        "\n",
        "\n",
        "    # Modifies `tokens_a` and `tokens_b` in place so that the total\n",
        "    # length is less than the specified length.\n",
        "    # Account for [CLS], [SEP], [SEP] with \"- 3\"\n",
        "    if len(tokens_a) > max_seq_length_src - 2:\n",
        "            tokens_a = tokens_a[0:(max_seq_length_src - 2)]\n",
        "    \n",
        "    if len(tokens_b) > max_seq_length_tgt - 2:\n",
        "            tokens_b = tokens_b[0:(max_seq_length_tgt - 2)]\n",
        "\n",
        "    \n",
        "    tokens_src = []\n",
        "    segment_ids_src = []\n",
        "    tokens_src.append(\"[CLS]\")\n",
        "    segment_ids_src.append(0)\n",
        "    for token in tokens_a:\n",
        "        tokens_src.append(token)\n",
        "        segment_ids_src.append(0)\n",
        "    tokens_src.append(\"[SEP]\")\n",
        "    segment_ids_src.append(0)\n",
        "  \n",
        "\n",
        "    tokens_tgt = []\n",
        "    segment_ids_tgt = []\n",
        "    tokens_tgt.append(\"[CLS]\")\n",
        "    #segment_ids_tgt.append(0)\n",
        "    for token in tokens_b:\n",
        "        tokens_tgt.append(token)\n",
        "        #segment_ids_tgt.append(0)\n",
        "    tokens_tgt.append(\"[SEP]\")\n",
        "    #segment_ids_tgt.append(0)\n",
        "\n",
        "    input_ids_src = tokenizer.convert_tokens_to_ids(tokens_src)\n",
        "   \n",
        "    \n",
        "\n",
        "    input_ids_tgt = tokenizer.convert_tokens_to_ids(tokens_tgt)\n",
        "    \n",
        "    #Adding begiining and end token\n",
        "    input_ids_tgt = input_ids_tgt[:-1] \n",
        "    \n",
        "    input_mask_src = [1] * len(input_ids_src)\n",
        "\n",
        "\n",
        "    input_mask_tgt = [1] * len(input_ids_tgt)\n",
        "    \n",
        "    labels_tgt = input_ids_tgt[1:]\n",
        "    \n",
        "    \n",
        "    labels_tgt.append(0)\n",
        "    \n",
        "    #print(len(input_ids_tgt))\n",
        "    #print(len(input_mask_tgt))\n",
        "    #print(len(labels_tgt))\n",
        "    #print(len(segment_ids_tgt))\n",
        "    \n",
        "    while len(input_ids_src) < max_seq_length_src:\n",
        "        input_ids_src.append(0)\n",
        "        input_mask_src.append(0)\n",
        "        segment_ids_src.append(0)\n",
        "\n",
        "    while len(input_ids_tgt) < max_seq_length_tgt:\n",
        "        input_ids_tgt.append(0)\n",
        "        input_mask_tgt.append(0)\n",
        "        segment_ids_tgt.append(0)\n",
        "        labels_tgt.append(0)\n",
        "\n",
        "    feature = InputFeatures( src_input_ids=input_ids_src,src_input_mask=input_mask_src,src_segment_ids=segment_ids_src,\n",
        "        tgt_input_ids=input_ids_tgt,tgt_input_mask=input_mask_tgt,tgt_labels=labels_tgt)\n",
        "\n",
        "    \n",
        "    return feature\n",
        "\n",
        "\n",
        "def file_based_input_fn_builder(input_file, max_seq_length_src,max_seq_length_tgt, is_training,\n",
        "                                drop_remainder, is_distributed=False):\n",
        "    \"\"\"Creates an `input_fn` closure to be passed to TPUEstimator.\"\"\"\n",
        "\n",
        "    name_to_features = {\n",
        "        \"src_input_ids\": tf.FixedLenFeature([max_seq_length_src], tf.int64),\n",
        "        \"src_input_mask\": tf.FixedLenFeature([max_seq_length_src], tf.int64),\n",
        "        \"src_segment_ids\": tf.FixedLenFeature([max_seq_length_src], tf.int64),\n",
        "        \"tgt_input_ids\": tf.FixedLenFeature([max_seq_length_tgt], tf.int64),\n",
        "        \"tgt_input_mask\": tf.FixedLenFeature([max_seq_length_tgt], tf.int64),\n",
        "        \"tgt_labels\" : tf.FixedLenFeature([max_seq_length_tgt], tf.int64),\n",
        "        \n",
        "        \n",
        "    }\n",
        "\n",
        "    def _decode_record(record, name_to_features):\n",
        "        \"\"\"Decodes a record to a TensorFlow example.\"\"\"\n",
        "        example = tf.parse_single_example(record, name_to_features)\n",
        "        print(example)\n",
        "        print(example.keys())\n",
        "\n",
        "        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.\n",
        "        # So cast all int64 to int32.\n",
        "        for name in list(example.keys()):\n",
        "            t = example[name]\n",
        "            if t.dtype == tf.int64:\n",
        "                t = tf.to_int32(t)\n",
        "            example[name] = t\n",
        "\n",
        "        return example\n",
        "\n",
        "    def input_fn(params):\n",
        "        \"\"\"The actual input function.\"\"\"\n",
        "        batch_size = params[\"batch_size\"]\n",
        "\n",
        "        # For training, we want a lot of parallel reading and shuffling.\n",
        "        # For eval, we want no shuffling and parallel reading doesn't matter.\n",
        "        d = tf.data.TFRecordDataset(input_file)\n",
        "        if is_training:\n",
        "\n",
        "            if is_distributed:\n",
        "                import horovod.tensorflow as hvd\n",
        "                tf.logging.info('distributed mode is enabled.'\n",
        "                                'size:{} rank:{}'.format(hvd.size(), hvd.rank()))\n",
        "                # https://github.com/uber/horovod/issues/223\n",
        "                d = d.shard(hvd.size(), hvd.rank())\n",
        "\n",
        "                d = d.repeat()\n",
        "                d = d.shuffle(buffer_size=100)\n",
        "                d = d.apply(\n",
        "                    tf.contrib.data.map_and_batch(\n",
        "                        lambda record: _decode_record(record, name_to_features),\n",
        "                        batch_size=batch_size//hvd.size(),\n",
        "                        drop_remainder=drop_remainder))\n",
        "            else:\n",
        "                tf.logging.info('distributed mode is not enabled.')\n",
        "                d = d.repeat()\n",
        "                d = d.shuffle(buffer_size=100)\n",
        "                d = d.apply(\n",
        "                    tf.contrib.data.map_and_batch(\n",
        "                        lambda record: _decode_record(record, name_to_features),\n",
        "                        batch_size=batch_size,\n",
        "                        drop_remainder=drop_remainder))\n",
        "\n",
        "        else:\n",
        "            d = d.apply(\n",
        "                tf.contrib.data.map_and_batch(\n",
        "                    lambda record: _decode_record(record, name_to_features),\n",
        "                    batch_size=batch_size,\n",
        "                    drop_remainder=drop_remainder))\n",
        "\n",
        "        return d\n",
        "    return input_fn\n",
        "  \n",
        "  \n",
        "def get_dataset(processor,\n",
        "                tokenizer,\n",
        "                data_dir,\n",
        "                max_seq_length_src,\n",
        "                max_seq_length_tgt,\n",
        "                batch_size,\n",
        "                mode,\n",
        "                output_dir,\n",
        "                is_distributed=False):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "        processor: Data Preprocessor, must have get_lables,\n",
        "            get_train/dev/test/examples methods defined.\n",
        "        tokenizer: The Sentence Tokenizer. Generally should be\n",
        "            SentencePiece Model.\n",
        "        data_dir: The input data directory.\n",
        "        max_seq_length: Max sequence length.\n",
        "        batch_size: mini-batch size.\n",
        "        model: `train`, `eval` or `test`.\n",
        "        output_dir: The directory to save the TFRecords in.\n",
        "    \"\"\"\n",
        "    #label_list = processor.get_labels()\n",
        "    if mode == 'train':\n",
        "        #train_examples = processor.get_train_examples(data_dir)\n",
        "        #train_file = os.path.join(output_dir, \"train.tf_record\")\n",
        "        train_file = \"gs://bert_summ/train.tf_record\"\n",
        "        #file_based_convert_examples_to_features(\n",
        "        #    train_examples, max_seq_length_src,max_seq_length_tgt,\n",
        "        #    tokenizer, train_file)\n",
        "        dataset = file_based_input_fn_builder(\n",
        "            input_file=train_file,\n",
        "            max_seq_length_src=max_seq_length_src,\n",
        "            max_seq_length_tgt =max_seq_length_tgt,\n",
        "            is_training=True,\n",
        "            drop_remainder=True,\n",
        "            is_distributed=is_distributed)({'batch_size': batch_size})\n",
        "    elif mode == 'eval':\n",
        "        #eval_examples = processor.get_dev_examples(data_dir)\n",
        "        #eval_file = os.path.join(output_dir, \"eval.tf_record\")\n",
        "        eval_file = \"gs://bert_summ/eval.tf_record\"\n",
        "        #file_based_convert_examples_to_features(\n",
        "        #    eval_examples, max_seq_length_src,max_seq_length_tgt,\n",
        "        #    tokenizer, eval_file)\n",
        "        dataset = file_based_input_fn_builder(\n",
        "            input_file=eval_file,\n",
        "            max_seq_length_src=max_seq_length_src,\n",
        "            max_seq_length_tgt =max_seq_length_tgt,\n",
        "            is_training=False,\n",
        "            drop_remainder=True,\n",
        "            is_distributed=is_distributed)({'batch_size': batch_size})\n",
        "    elif mode == 'test':\n",
        "      \n",
        "        #test_examples = processor.get_test_examples(data_dir)\n",
        "        #test_file = os.path.join(output_dir, \"predict.tf_record\")\n",
        "        test_file = \"gs://bert_summ/predict.tf_record\"\n",
        "        \n",
        "        #file_based_convert_examples_to_features(\n",
        "        #    test_examples, max_seq_length_src,max_seq_length_tgt,\n",
        "        #    tokenizer, test_file)\n",
        "        dataset = file_based_input_fn_builder(\n",
        "            input_file=test_file,\n",
        "            max_seq_length_src=max_seq_length_src,\n",
        "            max_seq_length_tgt =max_seq_length_tgt,\n",
        "            is_training=False,\n",
        "            drop_remainder=True,\n",
        "            is_distributed=is_distributed)({'batch_size': batch_size})\n",
        "    return dataset"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Nn4vhTvJjT0D",
        "colab_type": "code",
        "outputId": "f4080019-27c1-4f01-e498-a2b1d7600beb",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 227
        }
      },
      "cell_type": "code",
      "source": [
        "bert_config = model_utils.transform_bert_to_texar_config(\n",
        "            os.path.join(bert_pretrain_dir, 'bert_config.json'))\n",
        "\n",
        "\n",
        "\n",
        "tokenizer = tokenization.FullTokenizer(\n",
        "        vocab_file=os.path.join(bert_pretrain_dir, 'vocab.txt'),\n",
        "        do_lower_case=True)\n",
        "\n",
        "vocab_size = len(tokenizer.vocab)\n",
        "\n",
        "processor = CNNDailymail()\n",
        "train_dataset = get_dataset(processor,tokenizer,\"./\",max_seq_length_src,max_seq_length_tgt,4,'train',\"./\")\n",
        "eval_dataset = get_dataset(processor,tokenizer,\"./\",max_seq_length_src,max_seq_length_tgt,4,'eval',\"./\")\n",
        "test_dataset = get_dataset(processor,tokenizer,\"./\",max_seq_length_src,max_seq_length_tgt,4,'test',\"./\")\n",
        "#del processor"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:distributed mode is not enabled.\n",
            "WARNING:tensorflow:From <ipython-input-12-3918a62c9cd3>:297: map_and_batch (from tensorflow.contrib.data.python.ops.batching) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Use `tf.data.experimental.map_and_batch(...)`.\n",
            "{'src_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:0' shape=(512,) dtype=int64>, 'src_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:1' shape=(512,) dtype=int64>, 'src_segment_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:2' shape=(512,) dtype=int64>, 'tgt_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:3' shape=(400,) dtype=int64>, 'tgt_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:4' shape=(400,) dtype=int64>, 'tgt_labels': <tf.Tensor 'ParseSingleExample/ParseSingleExample:5' shape=(400,) dtype=int64>}\n",
            "dict_keys(['src_input_ids', 'src_input_mask', 'src_segment_ids', 'tgt_input_ids', 'tgt_input_mask', 'tgt_labels'])\n",
            "{'src_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:0' shape=(512,) dtype=int64>, 'src_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:1' shape=(512,) dtype=int64>, 'src_segment_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:2' shape=(512,) dtype=int64>, 'tgt_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:3' shape=(400,) dtype=int64>, 'tgt_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:4' shape=(400,) dtype=int64>, 'tgt_labels': <tf.Tensor 'ParseSingleExample/ParseSingleExample:5' shape=(400,) dtype=int64>}\n",
            "dict_keys(['src_input_ids', 'src_input_mask', 'src_segment_ids', 'tgt_input_ids', 'tgt_input_mask', 'tgt_labels'])\n",
            "{'src_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:0' shape=(512,) dtype=int64>, 'src_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:1' shape=(512,) dtype=int64>, 'src_segment_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:2' shape=(512,) dtype=int64>, 'tgt_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:3' shape=(400,) dtype=int64>, 'tgt_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:4' shape=(400,) dtype=int64>, 'tgt_labels': <tf.Tensor 'ParseSingleExample/ParseSingleExample:5' shape=(400,) dtype=int64>}\n",
            "dict_keys(['src_input_ids', 'src_input_mask', 'src_segment_ids', 'tgt_input_ids', 'tgt_input_mask', 'tgt_labels'])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "kDhy0XGlIcY2",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "del processor"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "CBXxrBteAuVj",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "424fc4a8-d889-412f-ae34-766b961ce7ab"
      },
      "cell_type": "code",
      "source": [
        "vocab_size"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "30522"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 15
        }
      ]
    },
    {
      "metadata": {
        "id": "lfw2JV11jsad",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#inputs to the model\n",
        "src_input_ids = tf.placeholder(tf.int64, shape=(None, None))\n",
        "src_segment_ids = tf.placeholder(tf.int64, shape=(None, None))\n",
        "tgt_input_ids = tf.placeholder(tf.int64, shape=(None, None))\n",
        "tgt_segment_ids = tf.placeholder(tf.int64, shape=(None, None))\n",
        "\n",
        "batch_size = tf.shape(src_input_ids)[0]\n",
        "\n",
        "src_input_length = tf.reduce_sum(1 - tf.to_int32(tf.equal(src_input_ids, 0)),\n",
        "                             axis=1)\n",
        "tgt_input_length = tf.reduce_sum(1 - tf.to_int32(tf.equal(src_input_ids, 0)),\n",
        "                             axis=1)\n",
        "\n",
        "labels = tf.placeholder(tf.int64, shape=(None, None))\n",
        "is_target = tf.to_float(tf.not_equal(labels, 0))\n",
        "\n",
        "\n",
        "global_step = tf.Variable(0, dtype=tf.int64, trainable=False)\n",
        "learning_rate = tf.placeholder(tf.float64, shape=(), name='lr')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "jTFde06_kACm",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#create the iterator \n",
        "iterator = tx.data.FeedableDataIterator({\n",
        "        'train': train_dataset, 'eval': eval_dataset, 'test': test_dataset})\n",
        "\n",
        "batch = iterator.get_next()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "XSrDO5YBkPYh",
        "colab_type": "code",
        "outputId": "a07623be-20ff-433c-a256-95f33b49b531",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "cell_type": "code",
      "source": [
        "#encoder Bert model\n",
        "print(\"Intializing the Bert Encoder Graph\")\n",
        "with tf.variable_scope('bert'):\n",
        "        embedder = tx.modules.WordEmbedder(\n",
        "            vocab_size=bert_config.vocab_size,\n",
        "            hparams=bert_config.embed)\n",
        "        word_embeds = embedder(src_input_ids)\n",
        "\n",
        "        # Creates segment embeddings for each type of tokens.\n",
        "        segment_embedder = tx.modules.WordEmbedder(\n",
        "            vocab_size=bert_config.type_vocab_size,\n",
        "            hparams=bert_config.segment_embed)\n",
        "        segment_embeds = segment_embedder(src_segment_ids)\n",
        "\n",
        "        input_embeds = word_embeds + segment_embeds\n",
        "\n",
        "        # The BERT model (a TransformerEncoder)\n",
        "        encoder = tx.modules.TransformerEncoder(hparams=bert_config.encoder)\n",
        "        encoder_output = encoder(input_embeds, src_input_length)\n",
        "        \n",
        "        # Builds layers for downstream classification, which is also initialized\n",
        "        # with BERT pre-trained checkpoint.\n",
        "        with tf.variable_scope(\"pooler\"):\n",
        "            # Uses the projection of the 1st-step hidden vector of BERT output\n",
        "            # as the representation of the sentence\n",
        "            bert_sent_hidden = tf.squeeze(encoder_output[:, 0:1, :], axis=1)\n",
        "            bert_sent_output = tf.layers.dense(\n",
        "                bert_sent_hidden, config_downstream.hidden_dim,\n",
        "                activation=tf.tanh)\n",
        "            output = tf.layers.dropout(\n",
        "                bert_sent_output, rate=0.1, training=tx.global_mode_train())\n",
        "\n",
        "\n",
        "print(\"loading the bert pretrained weights\")\n",
        "# Loads pretrained BERT model parameters\n",
        "init_checkpoint = os.path.join(bert_pretrain_dir, 'bert_model.ckpt')\n",
        "#init_checkpoint = \"gs://cloud-tpu-checkpoints/bert/uncased_L-12_H-768_A-12/bert_model.ckpt\"\n",
        "model_utils.init_bert_checkpoint(init_checkpoint)"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Intializing the Bert Encoder Graph\n",
            "loading the bert pretrained weights\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "C5m48bu5kVXm",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#decoder part and mle losss\n",
        "tgt_embedding = tf.concat(\n",
        "    [tf.zeros(shape=[1, embedder.dim]), embedder.embedding[1:, :]], axis=0)\n",
        "\n",
        "decoder = tx.modules.TransformerDecoder(embedding=tgt_embedding,\n",
        "                             hparams=dcoder_config)\n",
        "# For training\n",
        "outputs = decoder(\n",
        "    memory=encoder_output,\n",
        "    memory_sequence_length=src_input_length,\n",
        "    inputs=embedder(tgt_input_ids),\n",
        "    sequence_length=tgt_input_length,\n",
        "    decoding_strategy='train_greedy',\n",
        "    mode=tf.estimator.ModeKeys.TRAIN\n",
        ")\n",
        "\n",
        "mle_loss = transformer_utils.smoothing_cross_entropy(\n",
        "        outputs.logits, labels, vocab_size, loss_label_confidence)\n",
        "mle_loss = tf.reduce_sum(mle_loss * is_target) / tf.reduce_sum(is_target)\n",
        "\n",
        "train_op = tx.core.get_train_op(\n",
        "        mle_loss,\n",
        "        learning_rate=learning_rate,\n",
        "        global_step=global_step,\n",
        "        hparams=opt)\n",
        "\n",
        "tf.summary.scalar('lr', learning_rate)\n",
        "tf.summary.scalar('mle_loss', mle_loss)\n",
        "summary_merged = tf.summary.merge_all()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "sfDuR-SVkdhF",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#prediction \n",
        "start_tokens = tf.fill([tx.utils.get_batch_size(src_input_ids)],\n",
        "                       bos_token_id)\n",
        "predictions = decoder(\n",
        "    memory=encoder_output,\n",
        "    memory_sequence_length=src_input_length,\n",
        "    decoding_strategy='infer_greedy',\n",
        "    beam_width=beam_width,\n",
        "    alpha=alpha,\n",
        "    start_tokens=start_tokens,\n",
        "    end_token=eos_token_id,\n",
        "    max_decoding_length=400,\n",
        "    mode=tf.estimator.ModeKeys.PREDICT\n",
        ")\n",
        "if beam_width <= 1:\n",
        "    inferred_ids = predictions[0].sample_id\n",
        "else:\n",
        "    # Uses the best sample by beam search\n",
        "    inferred_ids = predictions['sample_id'][:, :, 0]\n",
        "\n",
        "\n",
        "saver = tf.train.Saver(max_to_keep=5)\n",
        "best_results = {'score': 0, 'epoch': -1}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "TCmgMIV6kzO4",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def _train_epoch(sess, epoch, step, smry_writer):\n",
        "        \n",
        "            \n",
        "        fetches = {\n",
        "            'step': global_step,\n",
        "            'train_op': train_op,\n",
        "            'smry': summary_merged,\n",
        "            'loss': mle_loss,\n",
        "        }\n",
        "\n",
        "        while True:\n",
        "            try:\n",
        "              feed_dict = {\n",
        "                iterator.handle: iterator.get_handle(sess, 'train'),\n",
        "                tx.global_mode(): tf.estimator.ModeKeys.TRAIN,\n",
        "              }\n",
        "              op = sess.run([batch],feed_dict)\n",
        "              feed_dict = {\n",
        "                   src_input_ids:op[0]['src_input_ids'],\n",
        "                   src_segment_ids : op[0]['src_segment_ids'],\n",
        "                   tgt_input_ids:op[0]['tgt_input_ids'],\n",
        "\n",
        "                   labels:op[0]['tgt_labels'],\n",
        "                   learning_rate: utils.get_lr(step, lr),\n",
        "                   tx.global_mode(): tf.estimator.ModeKeys.TRAIN\n",
        "                }\n",
        "\n",
        "\n",
        "              fetches_ = sess.run(fetches, feed_dict=feed_dict)\n",
        "              step, loss = fetches_['step'], fetches_['loss']\n",
        "              if step and step % display_steps == 0:\n",
        "                  logger.info('step: %d, loss: %.4f', step, loss)\n",
        "                  print('step: %d, loss: %.4f' % (step, loss))\n",
        "                  smry_writer.add_summary(fetches_['smry'], global_step=step)\n",
        "\n",
        "              if step and step % 1000 == 0:\n",
        "                  model_path = \"gs://bert_summ/models/model_\"+str(step)+\".ckpt\"\n",
        "                  logger.info('saving model to %s', model_path)\n",
        "                  print('saving model to %s' % model_path)\n",
        "                  saver.save(sess, model_path)\n",
        "              if step and step % eval_steps == 0:\n",
        "                  _eval_epoch(sess, epoch, mode='eval')\n",
        "            except tf.errors.OutOfRangeError:\n",
        "                break\n",
        "\n",
        "        return step"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "60_hbYdak5rd",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def _eval_epoch(sess, epoch, mode):\n",
        "\n",
        "        references, hypotheses = [], []\n",
        "        bsize = test_batch_size\n",
        "        fetches = {\n",
        "                'inferred_ids': inferred_ids,\n",
        "            }\n",
        "        bno=0\n",
        "        while True:\n",
        "            \n",
        "            #print(\"Temp\",temp)\n",
        "            try:\n",
        "              print(\"Batch\",bno)\n",
        "              feed_dict = {\n",
        "              iterator.handle: iterator.get_handle(sess, 'eval'),\n",
        "              tx.global_mode(): tf.estimator.ModeKeys.EVAL,\n",
        "              }\n",
        "              op = sess.run([batch],feed_dict)\n",
        "              feed_dict = {\n",
        "                   src_input_ids:op[0]['src_input_ids'],\n",
        "                   src_segment_ids : op[0]['src_segment_ids'],\n",
        "                   tx.global_mode(): tf.estimator.ModeKeys.EVAL\n",
        "              }\n",
        "              fetches_ = sess.run(fetches, feed_dict=feed_dict)\n",
        "              labels = op[0]['tgt_labels']\n",
        "              hypotheses.extend(h.tolist() for h in fetches_['inferred_ids'])\n",
        "              references.extend(r.tolist() for r in labels)\n",
        "              hypotheses = utils.list_strip_eos(hypotheses, eos_token_id)\n",
        "              references = utils.list_strip_eos(references, eos_token_id)\n",
        "              bno = bno+1\n",
        "              \n",
        "            except tf.errors.OutOfRangeError:\n",
        "                break\n",
        "\n",
        "\n",
        "        if mode == 'eval':\n",
        "            # Writes results to files to evaluate BLEU\n",
        "            # For 'eval' mode, the BLEU is based on token ids (rather than\n",
        "            # text tokens) and serves only as a surrogate metric to monitor\n",
        "            # the training process\n",
        "            #fname = os.path.join(model_dir, 'tmp.eval')\n",
        "            fname = \"./tmp.eval\"\n",
        "            #fname = \"gs://bert_summ/models/tmp.eval\"\n",
        "            hypotheses = tx.utils.str_join(hypotheses)\n",
        "            references = tx.utils.str_join(references)\n",
        "            hyp_fn, ref_fn = tx.utils.write_paired_text(\n",
        "                hypotheses, references, fname, mode='s')\n",
        "            eval_bleu = bleu_wrapper(ref_fn, hyp_fn, case_sensitive=True)\n",
        "            eval_bleu = 100. * eval_bleu\n",
        "            logger.info('epoch: %d, eval_bleu %.4f', epoch, eval_bleu)\n",
        "            print('epoch: %d, eval_bleu %.4f' % (epoch, eval_bleu))\n",
        "\n",
        "            if eval_bleu > best_results['score']:\n",
        "                logger.info('epoch: %d, best bleu: %.4f', epoch, eval_bleu)\n",
        "                best_results['score'] = eval_bleu\n",
        "                best_results['epoch'] = epoch\n",
        "                #model_path = os.path.join(model_dir, 'best-model.ckpt')\n",
        "                model_path = \"gs://bert_summ/models/best-model.ckpt\"\n",
        "                logger.info('saving model to %s', model_path)\n",
        "                print('saving model to %s' % model_path)\n",
        "                saver.save(sess, model_path)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "v9b6ghCsnT90",
        "colab_type": "code",
        "outputId": "bc4a0e8b-e9ad-408c-a92f-c6b4455b2d03",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 357
        }
      },
      "cell_type": "code",
      "source": [
        "#tx.utils.maybe_create_dir(model_dir)\n",
        "#logging_file = os.path.join(model_dir, 'logging.txt')\n",
        "\n",
        "model_dir = \"gs://bert_summ/models/\"\n",
        "logging_file= \"logging.txt\"\n",
        "logger = utils.get_logger(logging_file)\n",
        "with tf.Session() as sess:\n",
        "    sess.run(tf.global_variables_initializer())\n",
        "    sess.run(tf.local_variables_initializer())\n",
        "    sess.run(tf.tables_initializer())\n",
        "\n",
        "    smry_writer = tf.summary.FileWriter(model_dir, graph=sess.graph)\n",
        "\n",
        "    if run_mode == 'train_and_evaluate':\n",
        "        logger.info('Begin running with train_and_evaluate mode')\n",
        "\n",
        "        if tf.train.latest_checkpoint(model_dir) is not None:\n",
        "            logger.info('Restore latest checkpoint in %s' % model_dir)\n",
        "            saver.restore(sess, tf.train.latest_checkpoint(model_dir))\n",
        "        \n",
        "        iterator.initialize_dataset(sess)\n",
        "\n",
        "        step = 5000\n",
        "        for epoch in range(max_train_epoch):\n",
        "          iterator.restart_dataset(sess, 'train')\n",
        "          step = _train_epoch(sess, epoch, step, smry_writer)\n",
        "\n",
        "    elif run_mode == 'test':\n",
        "        logger.info('Begin running with test mode')\n",
        "\n",
        "        logger.info('Restore latest checkpoint in %s' % model_dir)\n",
        "        saver.restore(sess, tf.train.latest_checkpoint(model_dir))\n",
        "\n",
        "        _eval_epoch(sess, 0, mode='test')\n",
        "\n",
        "    else:\n",
        "        raise ValueError('Unknown mode: {}'.format(run_mode))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Restoring parameters from gs://bert_summ/models/model_5000.ckpt\n",
            "step: 5100, loss: 7.0686\n",
            "step: 5200, loss: 6.7414\n",
            "step: 5300, loss: 6.4176\n",
            "step: 5400, loss: 6.9609\n",
            "step: 5500, loss: 7.0777\n",
            "step: 5600, loss: 6.8462\n",
            "step: 5700, loss: 6.8764\n",
            "step: 5800, loss: 7.2216\n",
            "step: 5900, loss: 6.6034\n",
            "step: 6000, loss: 6.8505\n",
            "saving model to gs://bert_summ/models/model_6000.ckpt\n",
            "step: 6100, loss: 6.7107\n",
            "step: 6200, loss: 6.9797\n",
            "step: 6300, loss: 7.2937\n",
            "step: 6400, loss: 6.9824\n",
            "step: 6500, loss: 7.0897\n",
            "step: 6600, loss: 5.2173\n",
            "step: 6700, loss: 7.3187\n",
            "step: 6800, loss: 7.0490\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "L35eRRNKSoOV",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}